import json
import pytest
from pytest_httpserver import HTTPServer
from unittest.mock import patch
import requests
from string import Template

from tests.request_matching import (
    RequestMatcher,
    verify_request_made,
)
from tests.functional.app_config import AppConfig

pytestmark = pytest.mark.functional

path = "/api/conversation"
body = {
    "conversation_id": "123",
    "messages": [
        {"role": "user", "content": "Hello"},
        {"role": "assistant", "content": "Hi, how can I help?"},
        {"role": "user", "content": "What is the meaning of life?"},
    ],
}


@pytest.fixture(scope="function", autouse=True)
def setup_default_mocking(httpserver: HTTPServer, app_config: AppConfig):
    httpserver.expect_request(
        f"/openai/deployments/{app_config.get_from_json('AZURE_OPENAI_MODEL_INFO','model')}/chat/completions",
        method="POST",
    ).respond_with_data(
        Template(
            r"""data: {"id":"92f715be-cfc4-4ae6-80f8-c86b7955f6af","model":"$model","created":1712077271,"object":"extensions.chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","context":{"citations":[{"content":"document","title":"/documents/doc.pdf","url":{"id": "id", "source": "source", "title": "title", "chunk": 46, "chunk_id": null},"filepath":null,"chunk_id":"0"}],"intent":"[\"intent\"]"}},"end_turn":false,"finish_reason":null}]}

data: {"id":"92f715be-cfc4-4ae6-80f8-c86b7955f6af","model":"$model","created":1712077271,"object":"extensions.chat.completion.chunk","choices":[{"index":0,"delta":{"content":"42 is the meaning of life"},"end_turn":false,"finish_reason":null}],"system_fingerprint":"fp_68a7d165bf"}

data: {"id":"92f715be-cfc4-4ae6-80f8-c86b7955f6af","model":"$model","created":1712077271,"object":"extensio@ns.chat.completion.chunk","choices":[{"index":0,"delta":{},"end_turn":true,"finish_reason":"stop"}]}

data: [DONE]
"""
        ).substitute(model=app_config.get_from_json("AZURE_OPENAI_MODEL_INFO", "model"))
    )

    yield

    httpserver.check()


@patch(
    "backend.batch.utilities.search.azure_search_handler.AzureSearchHelper._index_not_exists"
)
@patch(
    "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default"
)
def test_azure_byod_responds_successfully_when_streaming(
    get_active_config_or_default_mock,
    index_not_exists_mock,
    app_url: str,
    app_config: AppConfig,
):
    get_active_config_or_default_mock.return_value.prompts.conversational_flow = "byod"
    index_not_exists_mock.return_value = False
    # when
    response = requests.post(f"{app_url}{path}", json=body)

    # then
    assert response.status_code == 200
    assert response.headers["Content-Type"] == "application/json-lines"

    response_lines = response.text.splitlines()
    assert len(response_lines) == 3

    final_response_json = json.loads(response_lines[-1])
    # Check only structure and key fields, not the exact content of citations which might contain dynamic SAS token
    assert "id" in final_response_json
    assert "model" in final_response_json
    assert "created" in final_response_json
    assert "object" in final_response_json
    assert "choices" in final_response_json
    assert len(final_response_json["choices"]) == 1
    assert "messages" in final_response_json["choices"][0]
    assert len(final_response_json["choices"][0]["messages"]) == 2

    # Check tool message
    tool_message = final_response_json["choices"][0]["messages"][0]
    assert tool_message["role"] == "tool"
    assert tool_message["end_turn"] is False
    assert "content" in tool_message

    # Parse citations from content
    tool_content = json.loads(tool_message["content"])
    assert "citations" in tool_content
    assert len(tool_content["citations"]) == 1
    citation = tool_content["citations"][0]
    assert "content" in citation
    assert "id" in citation
    assert "chunk_id" in citation
    assert "title" in citation
    assert "filepath" in citation
    assert "url" in citation

    # Check assistant message
    assistant_message = final_response_json["choices"][0]["messages"][1]
    assert assistant_message["role"] == "assistant"
    assert assistant_message["end_turn"] is True
    assert assistant_message["content"] == "42 is the meaning of life"


@patch(
    "backend.batch.utilities.search.azure_search_handler.AzureSearchHelper._index_not_exists"
)
@patch(
    "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default"
)
def test_post_makes_correct_call_to_azure_openai(
    get_active_config_or_default_mock,
    index_not_exists_mock,
    app_url: str,
    app_config: AppConfig,
    httpserver: HTTPServer,
):
    get_active_config_or_default_mock.return_value.prompts.use_on_your_data_format = (
        False
    )
    get_active_config_or_default_mock.return_value.prompts.conversational_flow = "byod"
    index_not_exists_mock.return_value = False
    # when
    requests.post(f"{app_url}{path}", json=body)

    # then
    verify_request_made(
        mock_httpserver=httpserver,
        request_matcher=RequestMatcher(
            path=f"/openai/deployments/{app_config.get_from_json('AZURE_OPENAI_MODEL_INFO','model')}/chat/completions",
            method="POST",
            query_string="api-version=2024-02-01",
            times=1,
            headers={
                "api-key": app_config.get("AZURE_OPENAI_API_KEY"),
            },
        ),
    )

    # Verify key parts of the request without being overly prescriptive about structure
    requests_log = httpserver.log
    found_matching_request = False

    for request_log in requests_log:
        request = request_log[0]
        if (request.path == f"/openai/deployments/{app_config.get_from_json('AZURE_OPENAI_MODEL_INFO','model')}/chat/completions" and request.method == "POST"):

            request_json = request.json

            # Check top-level fields
            assert "messages" in request_json
            assert "model" in request_json
            assert "temperature" in request_json
            assert "max_tokens" in request_json
            assert "top_p" in request_json
            assert "stream" in request_json
            assert "data_sources" in request_json

            # Check messages
            assert request_json["messages"] == body["messages"]

            # Check data_sources structure
            assert len(request_json["data_sources"]) == 1
            data_source = request_json["data_sources"][0]
            assert data_source["type"] == "azure_search"

            # Check data_source parameters
            parameters = data_source["parameters"]
            assert "endpoint" in parameters
            assert "index_name" in parameters
            assert "fields_mapping" in parameters
            assert "filter" in parameters
            assert "in_scope" in parameters
            assert "embedding_dependency" in parameters
            assert "query_type" in parameters
            assert "role_information" in parameters

            # Check fields_mapping
            fields_mapping = parameters["fields_mapping"]
            assert "content_fields" in fields_mapping
            assert "vector_fields" in fields_mapping
            assert "title_field" in fields_mapping
            assert "url_field" in fields_mapping
            assert "filepath_field" in fields_mapping

            found_matching_request = True
            break

    assert found_matching_request, "No matching request found"
